// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <melon/dynamic.h>
#include <pollux/common/time/cpu_wall_timer.h>
#include <pollux/exec/operator.h>

namespace kumo::pollux::exec {
    struct TaskStats;
} // namespace kumo::pollux::exec

namespace kumo::pollux::exec {
    /// Aggregated runtime statistics per plan node.
    ///
    /// Runtime statistics are collected on a per-operator instance basis. There can
    /// be multiple operator types and multiple instances of each operator type that
    /// correspond to a given plan node. For example, ProjectNode corresponds to
    /// a single operator type, FilterProject, but HashJoinNode corresponds to two
    /// operator types, HashProbe and HashBuild. Each operator type may have
    /// different runtime parallelism, e.g. there can be multiple instances of each
    /// operator type. Plan node statistics are calculated by adding up
    /// operator-level statistics for all corresponding operator instances.
    struct PlanNodeStats {
        explicit PlanNodeStats() = default;

        PlanNodeStats(const PlanNodeStats &) = delete;

        PlanNodeStats &operator=(const PlanNodeStats &) = delete;

        PlanNodeStats(PlanNodeStats &&) = default;

        PlanNodeStats &operator=(PlanNodeStats &&) = default;

        PlanNodeStats &operator+=(const PlanNodeStats &);

        /// Sum of input rows for all corresponding operators. Useful primarily for
        /// leaf plan nodes or plan nodes that correspond to a single operator type.
        uint64_t inputRows{0};

        /// Sum of input batches for all corresponding operators.
        vector_size_t inputVectors{0};

        /// Sum of input bytes for all corresponding operators.
        uint64_t inputBytes{0};

        /// Sum of raw input rows for all corresponding operators. Applies primarily
        /// to TableScan operator which reports rows before pushed down filter as raw
        /// input.
        uint64_t rawInputRows{0};

        /// Sum of raw input bytes for all corresponding operators.
        uint64_t rawInputBytes{0};

        /// Contains the dynamic filters stats if applied.
        DynamicFilterStats dynamicFilterStats;

        /// Sum of output rows for all corresponding operators. When
        /// plan node corresponds to multiple operator types, operators of only one of
        /// these types report non-zero output rows.
        uint64_t outputRows{0};

        /// Sum of output batches for all corresponding operators.
        vector_size_t outputVectors{0};

        /// Sum of output bytes for all corresponding operators.
        uint64_t outputBytes{0};

        // Sum of CPU, scheduled and wall times for isBLocked call for all
        // corresponding operators.
        CpuWallTiming isBlockedTiming;

        // Sum of CPU, scheduled and wall times for addInput call for all
        // corresponding operators.
        CpuWallTiming addInputTiming;

        // Sum of CPU, scheduled and wall times for noMoreInput call for all
        // corresponding operators.
        CpuWallTiming finishTiming;

        // Sum of CPU, scheduled and wall times for getOutput call for all
        // corresponding operators.
        CpuWallTiming getOutputTiming;

        /// Sum of CPU, scheduled and wall times for all corresponding operators. For
        /// each operator, timing of addInput, getOutput and finish calls are added
        /// up.
        CpuWallTiming cpuWallTiming;

        /// Sum of CPU, scheduled and wall times spent on background activities
        /// (activities that are not running on driver threads) for all corresponding
        /// operators.
        CpuWallTiming backgroundTiming;

        /// Sum of blocked wall time for all corresponding operators.
        uint64_t blockedWallNanos{0};

        /// Max of peak memory usage for all corresponding operators. Assumes that all
        /// operator instances were running concurrently.
        uint64_t peakMemoryBytes{0};

        uint64_t numMemoryAllocations{0};

        uint64_t physicalWrittenBytes{0};

        /// Operator-specific counters.
        melon::F14FastMap<std::string, RuntimeMetric> customStats;

        /// Breakdown of stats by operator type.
        melon::F14FastMap<std::string, std::unique_ptr<PlanNodeStats> > operatorStats;

        /// Number of drivers that executed the pipeline.
        int numDrivers{0};

        /// Number of total splits.
        int numSplits{0};

        /// Total bytes in memory for spilling
        uint64_t spilledInputBytes{0};

        /// Total bytes written for spilling.
        uint64_t spilledBytes{0};

        /// Total rows written for spilling.
        uint64_t spilledRows{0};

        /// Total spilled partitions.
        uint32_t spilledPartitions{0};

        /// Total spilled files.
        uint32_t spilledFiles{0};

        /// Add stats for a single operator instance.
        void add(const OperatorStats &stats);

        std::string toString(
            bool includeInputStats = false,
            bool includeRuntimeStats = false) const;

        bool isMultiOperatorTypeNode() const {
            return operatorStats.size() > 1;
        }

    private:
        void addTotals(const OperatorStats &stats);
    };

    melon::F14FastMap<core::PlanNodeId, PlanNodeStats> toPlanStats(
        const TaskStats &taskStats);

    melon::Dynamic toPlanStatsJson(const kumo::pollux::exec::TaskStats &stats);

    using PlanNodeAnnotation =
    std::function<std::string(const core::PlanNodeId &id)>;

    /// Returns human-friendly representation of the plan augmented with runtime
    /// statistics. The result has the same plan representation as in
    /// PlanNode::toString(true, true), but each plan node includes an additional
    /// line with runtime statistics. Plan nodes that correspond to multiple
    /// operator types, e.g. HashJoinNode, also include breakdown of runtime
    /// statistics per operator type.
    ///
    /// Note that input row counts and sizes are printed only for leaf plan nodes.
    ///
    /// @param includeCustomStats If true, prints operator-specific counters.
    /// @param annotation Function to extend plan printing with plugin, for
    /// example optimizer estimates next to execution stats.
    std::string printPlanWithStats(
        const core::PlanNode &plan,
        const TaskStats &taskStats,
        bool includeCustomStats = false,
        PlanNodeAnnotation annotation = nullptr);
} // namespace kumo::pollux::exec
